import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain.chat_models import ChatOpenAI
from langchain_community.agent_toolkits.load_tools import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from serpapi import BaiduSearch


# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
api_base = os.getenv("OPENAI_API_BASE")
serpapi_api_key = os.getenv("SERPAPI_API_KEY")
os.environ["OPENAI_API_KEY"] = api_key
os.environ["OPENAI_API_BASE"] = api_base
os.environ["SERPAPI_API_KEY"] = serpapi_api_key

query = "请问现任的美国总统是谁？他的年龄的平方是多少?"
# 使用百度进行搜索
# params = {
#   "engine": "baidu",
#   "q": query,
#   "api_key": serpapi_api_key
# }
#
# search = BaiduSearch(params)
# results = search.get_dict()
# print(results["organic_results"])

llm = ChatOpenAI(
    temperature=0,
    model="gpt-3.5-turbo",
)

tools = load_tools(
    ["serpapi","llm-math"],
    llm=llm
)

agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, # 0样本增强生成型
    verbose=True, # 打印日志
)
agent.run(query)

